"""
删除dump数据文件夹内所有npy文件名里的时间戳: python process_npy.py path/to/dump/rank_0/llama7b/0/0
打印npy信息: python process_npy.py path/to/dump/rank_0/llama7b/0/0/xxx.npy
比较两个npy文件夹: python process_npy.py path/to/npy_dir1 path/to/npy_dir2
比较两个npy文件: python process_npy.py path/to/npy_1.npy path/to/npy_2.npy
"""
import os
import sys
import shutil
import numpy as np
from glob import glob
import troubleshooter as ts

def compare(path1, path2):
    name_map = ts.migrator.get_name_map_list_by_number(path1, path2)
    ts.migrator.compare_npy_dir(path1, path2, name_map_list=name_map)

def find_diff_indices(arr1, arr2):
    if arr1.shape != arr2.shape:
        raise ValueError("两个数组必须具有相同的形状")
    diff_mask = arr1 != arr2
    diff_indices = np.where(diff_mask)
    return diff_mask, diff_indices

def remove_time_stamp(path):
    """删除dump数据文件夹内所有npy文件名里的时间戳"""
    file_dir, file_name = os.path.split(path)
    file_name_split = file_name.split(".")
    if "17" in file_name_split[-5]:
        del file_name_split[-5]
        new_file_name = ".".join(file_name_split)
        new_path = os.path.join(file_dir, new_file_name)
        os.rename(path, new_path)
        print("rename %s to %s", file_name, new_file_name)

def copy_dir(folder):
    folder_copy = folder + "_bak"
    if os.path.exists(folder_copy):
        shutil.rmtree(folder_copy)
    os.system(f"cp -r {folder} {folder_copy}")
    return folder_copy

if len(sys.argv) == 2:
    path1 = sys.argv[1]
    if os.path.isdir(path1):
        if path1.endswith("/"):
            path1 = path1[:-1]
        path1 = copy_dir(path1)
        npys = glob(os.path.join(path1, "*.npy"))
        for npy in npys:
            remove_time_stamp(npy)
    else:
        assert path1.endswith(".npy")
        npy = np.load(path1)
        print(npy)
elif len(sys.argv) >= 3:
    path1 = sys.argv[1]
    path2 = sys.argv[2]
    if os.path.isdir(path1) and os.path.isdir(path2):
        compare(path1, path2)
    elif os.path.isfile(path1) and os.path.isfile(path2):
        assert path1.endswith(".npy") and path2.endswith(".npy")
        npy_1 = np.load(path1)
        npy_2 = np.load(path2)
        diff_mask, diff_indices = find_diff_indices(npy_1, npy_2)
        print(diff_mask)
        print(diff_indices)
